import numpy as np
import networkx as nx
import pandas as pd
from IPython.display import clear_output
import matplotlib.pyplot as plt
import math
import random



def Add_Agents(y_0, State_Variables, Hyperparameters):
    
    """
    This function initializes agents, as well as their opinions and types
    
    Agents characteristics are stored in the dictionary State_Variables (key "Courteges")
    """
    
    #---------------------------------------------------------------------------
    
    M = Hyperparameters['M']
    m = Hyperparameters['m']
    N = Hyperparameters['N'] 
    N_bots = Hyperparameters['N_bots']
    #N_ordinary_agents = Hyperparameters['N_ordinary_agents']
    Init_Type = Hyperparameters['Init_Type']
    
    #---------------------------------------------------------------------------
    
    # Initialize opinions and types at the agent level
    
    if Init_Type == 'Random':  # opinions and types are endowed according to the specified joint distribution
        
        Courteges = []
        
        Y_0 = (y_0 * N).astype(int)  # obtain the distribution  from the starting point

        for i in range(m):
            for j in range(M):
                for k in range(Y_0[i, j]):

                    Courtege = [
                                i,   # opinion of agent
                                j,   # type of agent 
                                -1,  # this indicator makes sense only for bots
                               ]

                    Courteges.append(Courtege)
        
        random.shuffle(Courteges)  # random shuffle of agents' types and opinions 
                                   # comment this row to get assortative networks
        
    elif Init_Type == 'Custom':
        
        Courteges = Hyperparameters['Courteges']
            
    else:
        
        raise ValueError('Unknown initialization type!')
        
    
    Courteges = pd.DataFrame(Courteges, columns=['Opinion', 'Type', 'Target'])
                
    #---------------------------------------------------------------------------
    
    State_Variables['Courteges'] = Courteges
    
    return State_Variables



def Add_Bots(u, State_Variables, Hyperparameters):
    
    """
    This function initializes bots
    
    
    For each reconfiguration of influence structure, a new initialization of bots is launched 
    """
    
    #---------------------------------------------------------------------------
    
    M = Hyperparameters['M']
    m = Hyperparameters['m']
    N = Hyperparameters['N'] 
    N_bots = Hyperparameters['N_bots'] 
    Courteges = State_Variables['Courteges']
    
    #---------------------------------------------------------------------------
    
    U = (u * N).astype(int)
    
    #---------------------------------------------------------------------------
    
    Courteges_bots = []
    
    for i in range(m):
        
        for j in range(M):
            
            for k in range(U[i, j]):
                
                Courtege = [i, M, j]
                
                Courteges_bots.append(Courtege)
                
    Courteges_bots = pd.DataFrame(Courteges_bots, columns=['Opinion', 'Type', 'Target'])
                
    #---------------------------------------------------------------------------
    
    Courteges = pd.concat([Courteges[:N-N_bots], Courteges_bots], ignore_index=True)
    
    State_Variables['Courteges'] = Courteges
    
    return State_Variables



def Redefine_Network(State_Variables, Hyperparameters):
    
    """
    This function redefines the social network
    
    For each reconfiguration of control, a new social network is craeted 
    """
    
    M = Hyperparameters['M']
    m = Hyperparameters['m']
    N = Hyperparameters['N'] 
    N_bots = Hyperparameters['N_bots'] 
    N_ordinary_agents = Hyperparameters['N_ordinary_agents']
    Courteges = State_Variables['Courteges']
    
    Density = Hyperparameters['Density']
    
    Ordinary_Agents_Adj_Matrix = Hyperparameters['Ordinary_Agents_Adj_Matrix']
    
    #---------------------------------------------------------------------------
    
    # Adjacency matrix
    
    A = np.zeros((N, N))
    
    # Ordinary agents are connected by a complete graph
    
    #A[:(N-N_bots), :][:, :(N-N_bots)] = 1
    A[:N_ordinary_agents, :][:, :N_ordinary_agents] = Ordinary_Agents_Adj_Matrix
    
    # Define connections between ordinary agents and bots 
    
    for j in range(M):
        
        Target_Type = np.array(Courteges[Courteges['Type'] == j].index)
            
        Cohort = np.array(Courteges[Courteges['Target'] == j].index)
        
        for k in Target_Type:
            for s in Cohort:
                
                # The density of connections between ordinary agents and bots should exual to that of between ordinary agents themselves
                
                if random.random() <= Density:  
                
                    A[k, s] = 1
                    
                else:
                    
                    pass
        
    # Remove elements from the principal diagonal
    
    A = A - np.identity(N)
    
    A[A == (-1)] = 0
    
    #---------------------------------------------------------------------------
    
    # Create a graph object
    
    G = nx.from_numpy_array(A)
    
    State_Variables['G'] = G
    
    #---------------------------------------------------------------------------
    
    return State_Variables



def get_neighbors(agent, State_Variables, Hyperparameters):
    
    G = State_Variables['G']
    
    #---------------------------------------------------------------------------
    
    neighbors = []
    
    #---------------------------------------------------------------------------

    iterator = G.neighbors(agent)

    for neighbor in iterator:

        neighbors.append(neighbor)
        
    return neighbors



def Influence_Event(agent_i, agent_j, State_Variables, Hyperparameters):
    
    """
    This function imitates the contact between two agents. As a result, one agent updates their opinion
    according to the model protocol
    
    Input:
    agent_i         - index of the focal agent (influence object)
    agent_j         - index of their friend (influence source)
    State_Variables - the dictionary that stores stata variables (dynamically updates)
    Hyperparameters - the dictionary that stores hyperparameters of the experiment (constant)
    
    Ouptut:
    State_Variables - the dictionary that stores state variables (with an updated opinion of the focal agent)
    """
    
    # Take the table that stores agents' opinions and types
    
    Courteges = State_Variables['Courteges']
    
    # Retrieve opinions and types of intertacting agents
    
    Opinion_i = Courteges['Opinion'][agent_i]
    Opinion_j = Courteges['Opinion'][agent_j]
    
    Type_i = Courteges['Type'][agent_i]
    Type_j = Courteges['Type'][agent_j]
    
    # The corresponding transition matrix 
    
    Transition_Matrix = Hyperparameters['Transition_Matrices'][f'{Type_i}-{Type_j}']
    
    # Retrieve the probability distribution that describes the outcomes of the contact
    
    Distribution = Transition_Matrix[Opinion_i][Opinion_j]
    
    # Determine stochastically the outcome
    
    #print(len(Hyperparameters['Opinion_Alphabet']), len(Distribution))
    
    #print(Distribution)
    
    New_Opinion_i = np.random.choice(Hyperparameters['Opinion_Alphabet'], p=Distribution)
    
    #print(f'Old opinion: {Opinion_i}, new opinion: {New_Opinion_i}')
    
    #if Opinion_i != New_Opinion_i:
        
    #    print('Opinion has changed!')
        
    #else:
        
    #    pass
    
    # Insert new opinion
    
    #print(New_Opinion_i)
    
    Courteges.loc[agent_i, 'Opinion'] = New_Opinion_i
    
    # Update the dictionary of state variables 
    
    State_Variables['Courteges'] = Courteges
    
    return State_Variables



def Simulation_Run(u_t, State_Variables, Hyperparameters):
    
    """
    This function performes a simulation run of the agent-based model
    
    Input:
    u_t             - the list of control matrices
    Hyperparameters - the dictionary that stores hyperparameters of the experiment (constant)
    
    Ouptut:
    State_Variables - the dictionary that stores state variables 
    """
    
    #---------------------------------------------------------------------------
    
    T = Hyperparameters['T']  # the number of iterations (duration of the experiment)
    
    N = Hyperparameters['N']  # the number of agents
    N_ordinary_agents = Hyperparameters['N_ordinary_agents']  # the number of ordinary agents
    N_bots = Hyperparameters['N_bots']  # the number of bots
    
    Grid = Hyperparameters['Grid']
    
    #---------------------------------------------------------------------------
    
    Records = {}
    
    Records['States'] = []
    
    Records['Times'] = []
    
    #Records['Ass'] = []
    
    #---------------------------------------------------------------------------
    
    for t in range(1, T+1):
        

        #---------------------------------------------------------------------------
        
        if t in (Grid*N + 1):  # points of the grid - here we should redefine control and the network
        
            print(f'Iteration {t} of {T}')
            clear_output(wait=True)
            
            Records['States'].append(State_Variables['Courteges'][:(N-N_bots)].copy())
            Records['Times'].append(t)
            
            # Redefine control
            
            for k in range(len(Grid)):
                
                if t == (Grid[k]*N + 1):
                    
                    tau = k  # identify the index of the grid, which points out the control matrix that shoul be used
                    
                else:
                    
                    pass
                
            #print(tau)
            
            State_Variables = Add_Bots(u_t[tau], State_Variables, Hyperparameters)
            
            # Redefine graph
            
            State_Variables = Redefine_Network(State_Variables, Hyperparameters)
            
        else:
            
            pass
        
        #---------------------------------------------------------------------------
        

        agent_i = np.random.choice(N)  # select an agent
        
        if agent_i < (N - N_bots):  # if the chosen agent is an ordinary one

            neighbors = get_neighbors(agent_i, State_Variables, Hyperparameters)  # agent i's neighbors

            agent_j = np.random.choice(neighbors) # select one of their friends 

            State_Variables = Influence_Event(agent_i, agent_j, State_Variables, Hyperparameters)
            
        else:  # if the chosen agent is a bot
            
            pass  # do nothing
        
        """
        if t in (Grid*N + 1):
            
            #print(State_Variables['Courteges']['Opinion'].value_counts())
            #print(State_Variables['Courteges']['Opinion'][:(N-N_bots)].value_counts())
            
            Records['Opinions'].append(State_Variables['Courteges']['Opinion'][:(N-N_bots)].copy())
            Records['Times'].append(t)
            
        else:
            
            pass
            
        """
        
    Records['States'].append(State_Variables['Courteges'][:N_ordinary_agents])
            
    Records['Times'].append(t)
        
    return State_Variables, Records



def ABM_Dyn_System(y_0, u_t, Hyperparameters):
   
    """
    This function computes the state function using the agent-based modeling approach
    
    Input:
    y_0             - the starting point of the system (see formula (4) from the main manuscript)
    u_t             - the array of control matrices (control variable u (\tau) from the main manuscript) 
    Hyperparameters - hyperparameters
    
    Ouptut:
    y_t             - the array of state matrices (y (\tau) in the main manuscript)
    """

    #---------------------------------------------------------------------------
    
    m = Hyperparameters['m']
    M = Hyperparameters['M']
    N = Hyperparameters['N']  # the number of agents in the system
    
    #---------------------------------------------------------------------------
    
    # Expand the dictionary of hyperparameters
    
    N_bots = (u_t[0]*N).astype(int).sum()
    N_ordinary_agents = N - N_bots
    
    Hyperparameters['N_bots'] = N_bots
    Hyperparameters['N_ordinary_agents'] = N_ordinary_agents
    Hyperparameters['T'] = N * (Hyperparameters['Tau_1']-Hyperparameters['Tau_0'])
    
    Hyperparameters['Types'] = np.arange(M)
    Hyperparameters['Opinion_Alphabet'] = np.arange(m)
    
    #---------------------------------------------------------------------------
    
    State_Variables = {}

    State_Variables = Add_Agents(y_0, State_Variables, Hyperparameters)
    
    State_Variables, Records = Simulation_Run(u_t, State_Variables, Hyperparameters)
    
    #---------------------------------------------------------------------------
    
    Y_t = []  # this variable stand for the variable Y ( \tau ) (see the main manuscript)

    for snapshot in Records['States']:

        Y = np.zeros((m, M))

        for i in range(m):
            
            for j in range(M):
                
                #print(snapshot)

                Y[i, j] = sum((snapshot['Opinion'] == i) & (snapshot['Type'] == j) )

        Y_t.append(Y)

    
    # normalization
    
    y_t = [Y_t[i]/N for i in range(len(Y_t))]
    
    return y_t
